import cv2
import numpy as np
import time
import os
from sklearn import cluster


def layer(img: np.ndarray):
    height, width, channels = img.shape
    img = np.reshape(img, [-1, channels])

    km = cluster.KMeans(n_clusters=2)
    pred = km.fit_predict(img)
    pred = np.uint8(pred)
    # score = metrics.calinski_harabaz_score(pic, pred)
    # best_centers = [[round(e[0]), round(e[1]), round(e[2])] for e in km.cluster_centers_]
    best_centers = np.uint8(km.cluster_centers_)
    inds = np.sum(best_centers, 1)
    #
    # print(best_centers)
    # best_centers = best_centers[inds]
    # print(best_centers)
    #
    # pic = [best_centers[i] for i in pred]
    # pic = np.reshape(pic, [height, width, 3])
    # return pic
    if inds[0] > inds[1]:
        pred = 1 - pred  # [1 - i for i in pred]
    return np.reshape(pred, [height, width])


def make_dirs(path):
    pos = path.rfind(os.sep)
    if pos >= 0:
        path = path[:pos]
        os.makedirs(path, exist_ok=True)


colors = [
    [255, 0, 0],
    [0, 255, 0],
    [0, 0, 255],

    [255, 255, 0],
    [255, 0, 255],
    [0, 255, 255],
]


def count(binary_img, seps=None):
    assert len(seps) >= 2, 'seps contains only %d numbers, 2 needed at least' % len(seps)
    # pic = cv2.cvtColor(pic, cv2.COLOR_BGR2GRAY)
    if seps is None:
        seps = [0, min(binary_img.shape)]

    seps = [(seps[i], seps[i+1]) for i in range(len(seps)-1)]
    # ret, binary_img = cv2.threshold(binary_img, 128, 1, type=cv2.THRESH_BINARY)
    num, labels, stat, _ = cv2.connectedComponentsWithStats(binary_img, connectivity=8)

    areas = [0] * len(seps)
    # pic = np.uint8(pic * 255)
    # pic = cv2.cvtColor(pic, cv2.COLOR_GRAY2BGR)
    for x, y, w, h, n in stat:
        for i, (low, high) in enumerate(seps):
            if low <= max(w, h) < high:
                areas[i] += n
                # cv2.rectangle(pic, (x, y), (x+w, y+h), colors[i % len(colors)], 3)
                break
    # return areas, pic
    return np.array(areas)

# def main1():
#     pic = cv2.imread('../sample/caobo01.jpg')
#     layers = layer(pic)
#     layers = np.uint8(layers) * 255
#     pic1_BW = np.transpose([layers, layers, layers], [1, 2, 0])
#
#     cv2.imwrite('../sample/caobo01_result01_BW.jpg', pic1_BW)
#
#
# def main2():
#     pic1_BW = cv2.imread('../sample/caobo01_result01_BW.jpg')
#     seps = (100, 300, 500, 700, 900, min(pic1_BW.shape[:2]))
#     ns = count(pic1_BW, seps)
#     # cv2.imwrite('../sample/caobo01_result02_Red.jpg', pic)
#     ns = ns / np.sum(ns)
#     print(seps, sep='\t')
#     print(ns, sep='\t')
#     # cv2.imshow('aaa', cv2.resize(pic, (800, 800)))
#     # cv2.waitKey()


def get_distributes(img, seps, ratio=1.0):
    """
    获取图片中各种大小颗粒的占比
    :param img: 图片
    :param seps: 一个数值list或tuple，其中每个数值代表颗粒大小的一个界限。比如列表[100, 200, 300，500]表示三个区间：[100, 200),
                [200, 300)和[300,500)，其中每个区间包含下限不包含上限。根据用户的需要，数值的单位是公制（如毫米）、英制（如英寸）
                或者像素。本函数将根据ratio把数值转化为像素
    :param ratio: 像素对seps中数值的单位（比如毫米）的比例尺，比如120表示每毫米120个像素。
    :return: 返回一个numpy向量，含有len(seps)-1个比例，表示seps中每个区间对应的颗粒在所有颗粒中所占比例。所有占比之和等于1.0。
    """
    img = layer(img)
    seps = [round(e * ratio) for e in seps]
    ns = count(img, seps)
    ns = ns ** 1.5
    return ns / np.sum(ns)


from util.func_util import print_expired

@print_expired
def main3():
    pic = cv2.imread('../sample/caobo01.jpg')
    h, w, _ = pic.shape
    print('h, w = ', h, w)

    size = 1200
    # r = size / max(w, h)
    r = 0.5
    if h > size or w > size:
        pic = cv2.resize(pic, (int(w * r), int(h * r)))
    # r = 1.0
    seps = np.array([100, 200, 300, 500, 700, 1000, 1500]) * r
    print(get_distributes(pic, seps))


if __name__ == '__main__':
    main3()